import socket
import shutil
import traceback
import os
import sys
import subprocess
import torch
import sys
import textwrap

import re
def print_boxed_description(description, char="*"):
    """Prints a boxed description, adapting to terminal width and truncating if needed."""

    terminal_width = shutil.get_terminal_size().columns

    # Remove ANSI escape codes for length calculation
    ansi_escape = re.compile(r'\x1b$$[0-9;]*[mK]')
    description_stripped = ansi_escape.sub('', description)

    inner_width = terminal_width 

    if len(description_stripped) > inner_width:
        wrapped_text = textwrap.fill(description_stripped, width=inner_width)
        lines = wrapped_text.splitlines()
        if len(lines) > 2:
            description = "\n".join(lines[:2]) + " ..."
        else:
            description = wrapped_text  # Use the wrapped text directly
    else:
        description = description.ljust(inner_width)  # Left justify within inner width

    box = [char * terminal_width]
    for line in description.splitlines():
        line = line.rstrip()  # Remove trailing spaces
        box.append(f"> {line.ljust(inner_width)} ")
    box.append(char * terminal_width)

    print("\n".join(box))


def is_port_open(port, host='127.0.0.1'):
    """check if the specified port is occupied"""
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(0.1)  
    result = sock.connect_ex((host, port))  # try to connect to the port
    sock.close()
    return result == 0  # if the port is occupied, the result will be 0

def get_free_port():
    """get a free port"""
    port = 6007
    while port < 7000:
        if is_port_open(port) == False:
            return port
        port += 1
    raise Exception(f"No free port found between 6000 and 7000")

def copy_local_modules(target_directory, stack_index:int = 0):
    """copy local modules to target directory
    
    Args:
        target_directory (str): target directory
        stack_index (int, optional): stack index. Defaults to 0.
    """
    def check_imported_modules():
        def is_local_module(module):
            """Check if the module is locally imported"""
            if hasattr(module, '__file__') and isinstance(module.__file__, str):
                module_path = module.__file__
                return module_path.startswith(os.getcwd()) 
            return False
        local_modules = []
        # Iterate over imported modules
        for module_name, module in sys.modules.items():
            # Process only if the module name format is valid
            if module_name.startswith('_') is False and module_name.isidentifier():
                if is_local_module(module):
                    local_modules.append(module_name)
        
        return local_modules
    
    local_mods = check_imported_modules()
    """Copy the files of local modules to the specified directory"""
    if not os.path.exists(target_directory):
        os.makedirs(target_directory)

    for module_name in local_mods:
        module = sys.modules[module_name]
        if hasattr(module, '__file__'):
            module_file = module.__file__
            module_folder, _ = os.path.split(module_file)
            # Copy the folder
            shutil.copytree(module_folder, os.path.join(target_directory, os.path.basename(module_folder)), dirs_exist_ok=True)
            print(f"Copied {module_folder} to {target_directory}")
    
    stack = traceback.extract_stack()
    # -2 is the stack frame of the current function call, -1 is the stack frame of the current function
    caller_frame = stack[-2 - stack_index]
    # caller_frame[0] is the file path, caller_frame[1] is the line number, etc.
    caller_file_path = caller_frame[0]

    shutil.copy(caller_file_path, target_directory)

if __name__ == '__main__':
    import sys
    port_to_check = int(sys.argv[1]) if len(sys.argv) > 1 else 6007
    if is_port_open(port_to_check):
        print(f"the port {port_to_check} is occupied")
    else:
        print(f"the port {port_to_check} is available")
